package com.lighting.common.log;


import brave.Span;
import brave.Tracer;
import brave.Tracing;
import brave.propagation.TraceContext;
import org.slf4j.MDC;

/**
 * @author zhanghanjun
 * @date 2022-07-01 11:29
 */
public class TraceUtils {

    public static TraceBuilder getTraceBuilder() {
        Span span = Tracing.current().tracer().currentSpan();
        if (span == null) {
            return null;
        }
        TraceContext context = span.context();
        if (context == null) {
            return null;
        }
        TraceBuilder traceBuilder = new TraceBuilder();
        traceBuilder.setTraceId(context.traceId());
        traceBuilder.setSpanId(context.spanId());
        traceBuilder.setParentId(context.parentIdAsLong());
        traceBuilder.setTraceIdHigh(context.traceIdHigh());
        traceBuilder.setLocalRootId(context.localRootId());
        traceBuilder.setExtraList(context.extra());
        return traceBuilder;
    }

    public static void withBuilderInTrace(TraceBuilder traceBuilder) {
        if (traceBuilder == null) {
            return;
        }

        Tracer tracer = Tracing.current().tracer();
        TraceContext.Builder builder = tracer.nextSpan().context().toBuilder();
        builder.traceId(traceBuilder.getTraceId());
        builder.spanId(traceBuilder.getSpanId());
        builder.parentId(traceBuilder.getParentId());
        builder.traceIdHigh(traceBuilder.getTraceIdHigh());
        traceBuilder.getExtraList().stream().forEach(extra -> builder.addExtra(extra));

        TraceContext traceContext = builder.build();
        if (traceContext.traceIdString().equals(MDC.get(TraceConstant.MDC_TRACE_ID))) {
            return;
        }

        Span span = tracer.toSpan(traceContext);
        tracer.withSpanInScope(span.start());
    }

}
